#include "kernel_operator.h"

using namespace AscendC;

template <typename T>
class KernelOptensor
{
    const uint32_t BUFFER_NUM = 2;

public:
    __aicore__ inline KernelOptensor() {}
    __aicore__ inline void Init(TPipe *pipeIn, GM_ADDR x1, GM_ADDR x2, GM_ADDR y,
                                uint32_t coreDataNum, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t tailDataNum)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->coreDataNum = coreDataNum;
        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = tailDataNum;

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, this->coreDataNum);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, this->coreDataNum);
        yGm.SetGlobalBuffer((__gm__ T *)y, this->coreDataNum);

        pipe = pipeIn;
        pipe->InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(T));
        if constexpr (std::is_same_v<T, half> || std::is_same_v<T, bfloat16_t>)
        {
            pipe->InitBuffer(tmp1, this->tileDataNum * sizeof(float));
            pipe->InitBuffer(tmp2, this->tileDataNum * sizeof(float));
        }
    }
    __aicore__ inline void Process()
    {
        uint32_t loopCount = this->tileNum;
        this->processDataNum = this->tileDataNum;
        for (uint32_t i = 0; i < loopCount; i++)
        {
            if (i == loopCount - 1)
            {
                this->processDataNum = this->tailDataNum;
            }
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(uint32_t progress)
    {
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopy(x1Local, x1Gm[progress * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[progress * this->tileDataNum], this->processDataNum);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    __aicore__ inline void Compute(uint32_t progress)
    {
        LocalTensor<T> x1Local = inQueueX1.DeQue<T>();
        LocalTensor<T> x2Local = inQueueX2.DeQue<T>();
        LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
        if constexpr (std::is_same_v<T, half> || std::is_same_v<T, bfloat16_t>){
            auto p1 = tmp1.Get<float>();
            auto p2 = tmp2.Get<float>();
            Cast(p1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Cast(p2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
            Ln(p1, p1, this->processDataNum);
            Mul(p1, p1, p2, this->processDataNum);
            Exp(p1, p1, this->processDataNum);
            Cast(yLocal, p1, RoundMode::CAST_ROUND, this->processDataNum);
        }else{
            Ln(x1Local, x1Local, this->processDataNum);
            Mul(x1Local, x1Local, x2Local, this->processDataNum);
            Exp(yLocal, x1Local, this->processDataNum);
        }

        outQueueY.EnQue(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
    }
    __aicore__ inline void CopyOut(uint32_t progress)
    {
        LocalTensor<T> yLocal = outQueueY.DeQue<T>();
        DataCopy(yGm[progress * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe *pipe;
    TQue<QuePosition::VECIN, 1> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, 1> outQueueY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;

    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;

    uint32_t coreDataNum;
    uint32_t tileNum;
    uint32_t tileDataNum;
    uint32_t tailDataNum;
    uint32_t processDataNum;
};
